import seaborn as sns

# 权重数据（高温态 vs 低温态）
np.random.seed(42)
high_T_weights = np.random.dirichlet(np.ones(100)*0.5, 1)[0]  # 高温态平坦分布
low_T_weights = np.random.dirichlet(np.ones(100)*10, 1)[0]    # 低温态尖锐分布

# 绘制热力图对比
fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# 高温态权重分布
sns.heatmap(high_T_weights.reshape(10,10), ax=axes[0], cmap="YlGnBu", 
            cbar_kws={'label': 'Attention Weight'}, vmin=0, vmax=0.05)
axes[0].set_title(r'High Temperature ($T^{(t)}=0.7$)', fontsize=12)
axes[0].set_xlabel('Device Index', fontsize=10)
axes[0].set_ylabel('Device Index', fontsize=10)

# 低温态权重分布
sns.heatmap(low_T_weights.reshape(10,10), ax=axes[1], cmap="YlGnBu", 
            cbar_kws={'label': 'Attention Weight'}, vmin=0, vmax=0.05)
axes[1].set_title(r'Low Temperature ($T^{(t)}=0.1$)', fontsize=12)
axes[1].set_xlabel('Device Index', fontsize=10)

plt.suptitle('Layer-wise Attention Weight Distribution', fontsize=14)
plt.savefig('attention_weight_heatmap.png', bbox_inches='tight')
